import os.path
from PIL import Image
import numpy as np
import cv2


def divide(filename):
    # Open the image
    img = Image.open(filename)
    
    # Convert PIL image to OpenCV format (BGR)
    img_cv = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)
    
    # Get image dimensions
    height, width = img_cv.shape[:2]
    
    # Since it's black background with white text, we sum along the vertical axis
    # to get a horizontal projection of the text
    gray = cv2.cvtColor(img_cv, cv2.COLOR_BGR2GRAY)
    horizontal_projection = np.sum(gray, axis=0)
    
    # Find the approximate centers of the 5 columns
    # We divide the width into 5 parts and find the center of text in each part
    part_width = width // 5
    column_centers = []
    
    for i in range(5):
        start_x = i * part_width
        end_x = min((i + 1) * part_width, width)
        
        # Get the segment for this column
        segment = horizontal_projection[start_x:end_x]
        
        # Find the center of mass of the text in this segment
        if np.sum(segment) > 0:  # If there's text in this segment
            # Calculate weighted center
            indices = np.arange(len(segment))
            center = np.average(indices, weights=segment+1e-10)  # Add small value to avoid zero weights
            column_centers.append(int(start_x + center))
        else:
            # If no text found, use geometric center
            column_centers.append(start_x + part_width // 2)
    
    # Now determine boundaries between columns
    column_boundaries = [0]  # Start at the very left
    
    # Calculate boundaries between adjacent columns
    for i in range(len(column_centers) - 1):
        boundary = (column_centers[i] + column_centers[i + 1]) // 2
        column_boundaries.append(boundary)
    
    column_boundaries.append(width)  # End at the very right
    
    # Remove the manual boundary adjustment as we now use sufficient overlap
    # and rely on the intrusion removal function for clean separation
    adjusted_boundaries = column_boundaries
    
    # Extract the base filename without extension
    base_name = os.path.splitext(os.path.basename(filename))[0]
    directory = os.path.dirname(filename)
    
    # Split the image into 5 columns and save each
    # For Chinese texts, columns are read from right to left
    for i in range(5):
        # Calculate position from right to left
        col_index = 4 - i
        
        left = adjusted_boundaries[col_index]
        right = adjusted_boundaries[col_index + 1]
        
        # Crop the column with increased overlap to better preserve character edges
        overlap = 20
        crop_left = max(0, left - overlap)
        crop_right = min(width, right + overlap)
        
        # Crop the column
        column_img = img_cv[:, crop_left:crop_right]
        
        # Adjust boundaries for processing
        process_left = left - crop_left
        process_right = right - crop_left
        
        # Process the column to remove intrusion from adjacent columns
        processed_column = remove_intrusion(column_img, process_left, process_right)
        
        # Convert back to PIL Image for saving
        processed_pil = Image.fromarray(cv2.cvtColor(processed_column, cv2.COLOR_BGR2RGB))
        
        # Create output filename with "-" prefix and reversed numbering
        output_filename = os.path.join(directory, f"{base_name}-{i+1}.jpg")
        processed_pil.save(output_filename)


def remove_intrusion(img, left_boundary, right_boundary):
    """
    Remove intrusion from adjacent columns using the specified approach:
    1. Preprocess: grayscale -> adaptive threshold -> morphological opening
    2. Divide into regions: 10% edge regions, 80% central region
    3. Identify intrusion: extract white connected components in edge regions,
       filter out components with area < 50 or connected to central region
    4. Fill: change identified intrusion regions to black
    """
    # Step 1: Preprocessing
    # Convert to grayscale
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    
    # Apply adaptive threshold to get binary image
    binary = cv2.adaptiveThreshold(gray, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, 
                                   cv2.THRESH_BINARY, 11, 2)
    
    # Apply morphological opening to remove noise
    kernel = np.ones((3,3), np.uint8)
    cleaned = cv2.morphologyEx(binary, cv2.MORPH_OPEN, kernel)
    
    # Step 2: Divide into regions
    height, width = cleaned.shape
    edge_width = max(1, width // 10)  # 10% of width, at least 1 pixel
    
    # Create a copy of the original image for result
    result_img = img.copy()
    
    # Process left edge if it exists
    if edge_width > 0:
        # Define the left edge region (area that might contain intrusion from left neighbor)
        left_edge_start = max(0, left_boundary - edge_width)
        left_edge_end = left_boundary
        
        # Make sure we have a valid region
        if left_edge_end > left_edge_start:
            # Step 3: Identify intrusion in left edge
            # Find connected components in the left edge
            num_labels, labels, stats, centroids = cv2.connectedComponentsWithStats(
                cleaned[:, left_edge_start:left_edge_end], connectivity=8)
            
            # Process each component in the left edge
            for label in range(1, num_labels):  # Skip background (label 0)
                area = stats[label, cv2.CC_STAT_AREA]
                # Filter out small components (area < 50)
                if area >= 50:
                    # Create mask for this component
                    component_mask = (labels == label).astype(np.uint8) * 255
                    # Fill the intrusion area with black in the result image
                    result_img[:, left_edge_start:left_edge_end][component_mask == 255] = [0, 0, 0]
        
        # Define the right edge region (area that might contain intrusion from right neighbor)
        right_edge_start = right_boundary
        right_edge_end = min(width, right_boundary + edge_width)
        
        # Make sure we have a valid region
        if right_edge_end > right_edge_start:
            # Step 3: Identify intrusion in right edge
            # Find connected components in the right edge
            num_labels, labels, stats, centroids = cv2.connectedComponentsWithStats(
                cleaned[:, right_edge_start:right_edge_end], connectivity=8)
            
            # Process each component in the right edge
            for label in range(1, num_labels):  # Skip background (label 0)
                area = stats[label, cv2.CC_STAT_AREA]
                # Filter out small components (area < 50)
                if area >= 50:
                    # Create mask for this component
                    component_mask = (labels == label).astype(np.uint8) * 255
                    # Fill the intrusion area with black in the result image
                    result_img[:, right_edge_start:right_edge_end][component_mask == 255] = [0, 0, 0]
    
    return result_img


def main():
    image_name = "01.jpg"
    # Fix the path to correctly point to the raw directory
    filename = os.path.join(os.path.dirname(__file__), "..", "raw", image_name)
    # Normalize the path
    filename = os.path.normpath(filename)
    divide(filename)


if __name__ == "__main__":
    main()